#include "../DWMRICore/Fiber.h"
#include "../DWMRICore/Scalar.h"
#include "../DWMRICore/DTI.h"

#include "LeastSquare.h"
 
#include <teem\ell.h>
 
#include <omp.h>
 
char *fiber_pathname = NULL;
char *dst_pathname = NULL;
 
char **confPathname = NULL;
char **vertPathname = NULL;
int w, h, d, scale;
float stepSize, maxLength;
 
float fiber_length = 0.0f;
 
struct Euclidean{
    Seed seed;
    float parameters[9];
};
 
 
void PrintUsage()
{
    printf("Usage: ComputeEuclidean -fiber <configure_file> -dst <dst_file> -length <fiber_length>\n");
}
 
int ParseArguments(const int argc, char *argv[])
{
    for (int i = 1; i < argc; ++i) {
        if (strcmp(argv[i], "-fiber") == 0) {
            fiber_pathname = argv[++i];
        } else if (strcmp(argv[i], "-dst") == 0) {
            dst_pathname = argv[++i];
        } else if (strcmp(argv[i], "-length") == 0) {
            fiber_length = (float)atof(argv[++i]);
        } else {
            return -1;
        }
    }
 
    if (fiber_pathname == NULL || dst_pathname == NULL || fiber_length < 0.0001f)
        return -2;
 
    return 0;
}
 
void MatrixSquareRoot(const float *src, float *dst)
{
    double m[9], eval[3], evec[9];
    m[0] = src[0];  m[1] = src[1];  m[2] = src[2];
    m[3] = src[1];  m[4] = src[3];  m[5] = src[4];
    m[6] = src[2];  m[7] = src[4];  m[8] = src[5];
 
    ell_3m_eigensolve_d(eval, evec, m, AIR_TRUE);
 
    for (int i = 0; i < 3; ++i) {
        if (eval[i] > 0.0)
            eval[i] = sqrt(eval[i]);
        else
            eval[i] = 0.0f;
    }
 
    dst[0] = (float)(eval[0] * evec[0] * evec[0] + eval[1] * evec[3] * evec[3] + eval[2] * evec[6] * evec[6]);
    dst[1] = (float)(eval[0] * evec[0] * evec[1] + eval[1] * evec[3] * evec[4] + eval[2] * evec[6] * evec[7]);
    dst[2] = (float)(eval[0] * evec[0] * evec[2] + eval[1] * evec[3] * evec[5] + eval[2] * evec[6] * evec[8]);
    dst[3] = (float)(eval[0] * evec[1] * evec[1] + eval[1] * evec[4] * evec[4] + eval[2] * evec[7] * evec[7]);
    dst[4] = (float)(eval[0] * evec[1] * evec[2] + eval[1] * evec[4] * evec[5] + eval[2] * evec[7] * evec[8]);
    dst[5] = (float)(eval[0] * evec[2] * evec[2] + eval[1] * evec[5] * evec[5] + eval[2] * evec[8] * evec[8]);
}
 
void EuclideanParameterization(CFiber *fibers, Euclidean *parameters, const int count)
{
    uint length = uint(fiber_length / stepSize);
    for (int i = 0; i < count; ++i) {
        uint F = fibers[i].m_fCount > length ? length : fibers[i].m_fCount;
        uint B = fibers[i].m_bCount > length ? length : fibers[i].m_bCount;
 
        parameters[i].seed.pos = fibers[i].m_seed.pos;
        parameters[i].seed.dir = fibers[i].m_seed.dir;
 
        if (F + B <= 2)
            continue;
 
        /* mean vector */
        CVector3F mean(0.0f, 0.0f, 0.0f);
        for (uint k = F-1; k > 0; --k) {
            CVector3F temp = fibers[i].m_pF[k-1] - fibers[i].m_pF[k];
            mean = mean + temp;
        }
        for (uint k = 0; k < B-1; ++k) {
            CVector3F temp = fibers[i].m_pB[k+1] - fibers[i].m_pB[k];
            mean = mean + temp;
        }
		//mean.Normalize();
        mean = mean / float(F + B - 2);
         
        /* covariance matrix */
        float cov[6] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
        for (uint k = F-1; k > 0; --k) {
            CVector3F temp = fibers[i].m_pF[k-1] - fibers[i].m_pF[k];
			//temp.Normalize();
            CVector3F d = temp - mean;
            cov[0] = d.m_x * d.m_x; cov[1] = d.m_x * d.m_y; cov[2] = d.m_x * d.m_z;
            cov[3] = d.m_y * d.m_y; cov[4] = d.m_y * d.m_z; cov[5] = d.m_z * d.m_z;
        }
        for (uint k = 0; k < B-1; ++k) {
            CVector3F temp = fibers[i].m_pB[k+1] - fibers[i].m_pB[k];
			//temp.Normalize();
            CVector3F d = temp - mean;
            cov[0] = d.m_x * d.m_x; cov[1] = d.m_x * d.m_y; cov[2] = d.m_x * d.m_z;
            cov[3] = d.m_y * d.m_y; cov[4] = d.m_y * d.m_z; cov[5] = d.m_z * d.m_z;
        }
        cov[0] /= float(F + B - 2); cov[1] /= float(F + B - 2); cov[2] /= float(F + B - 2);
        cov[3] /= float(F + B - 2); cov[4] /= float(F + B - 2); cov[5] /= float(F + B - 2);
 
        parameters[i].parameters[0] = mean.m_x;
        parameters[i].parameters[1] = mean.m_y;
        parameters[i].parameters[2] = mean.m_z;
        MatrixSquareRoot(cov, &(parameters[i].parameters[3]));
    }
}
 
void ReverseFiber(Euclidean *fiber)
{
    fiber->parameters[0] = -fiber->parameters[0];
    fiber->parameters[1] = -fiber->parameters[1];
    fiber->parameters[2] = -fiber->parameters[2];
}
 
/* estimate a function for each parameter, and compute the edge strength from it */
float ComputeEdgeStrengthEuclidean(Euclidean fibers[27], const int m)
{
    /* matrix A is same for all parameters */
	double *A = new double[m*4];
	for (int i = 0; i < m; ++i) { 
		A[i*4+0] = (fibers[i].seed.pos.m_x - fibers[0].seed.pos.m_x) * (double)scale;
		A[i*4+1] = (fibers[i].seed.pos.m_y - fibers[0].seed.pos.m_y) * (double)scale;
		A[i*4+2] = (fibers[i].seed.pos.m_z - fibers[0].seed.pos.m_z) * (double)scale;
		A[i*4+3] = 1.0;
	}

	/* compute pseudoinverse of matrix A */
	double *pinvA = new double[4*m];
	PseudoInverse4(A, pinvA, m, 4);

	float strength = 0.0f;
 
	float *x = new float[3]; 
    float *y = new float[m];

    for (int i = 0; i < 9; ++i) {
          for (int k = 0; k < m; ++k) {
             y[k] = fibers[k].parameters[i];
         }
 
         memset(x, 0, sizeof(float)*3);
 
         LeastSquareFast4(pinvA, x, y, m, 4);
 
         strength += x[0] * x[0] + x[1] * x[1] + x[2] * x[2];
    }

	SafeDeleteArray(A);
	SafeDeleteArray(pinvA);
	SafeDeleteArray(x);
	SafeDeleteArray(y);
 
    return sqrt(strength);
}
 
int main(int argc, char *argv[])
{
    if (ParseArguments(argc, argv) != 0) {
        PrintUsage();
        return 0;
    }
 
    /* read the configure file */
    ReadConfigureFile(fiber_pathname, &confPathname, &vertPathname, w, h, d, scale, stepSize, maxLength);
 
    float length = fiber_length > maxLength ? maxLength : fiber_length;
    int cw, ch;
 
    /* compute the edge strength for each voxel */
    float *edgeStrength = new float[w*h*d];
    memset(edgeStrength, 0, sizeof(float)*w*h*d);
 
    CFiber *fibers = NULL;
 
    Euclidean **euclidean = new Euclidean*[3];
    euclidean[0] = new Euclidean[w*h];
    euclidean[1] = new Euclidean[w*h];
    euclidean[2] = new Euclidean[w*h];
    memset(euclidean[0], 0, sizeof(Euclidean)*w*h);
    memset(euclidean[1], 0, sizeof(Euclidean)*w*h);
    memset(euclidean[2], 0, sizeof(Euclidean)*w*h);
 
    int offset[9];
    offset[ 0] = 0;     offset[ 1] = -1;    offset[ 2] = 1;
    offset[ 3] = -w-1;  offset[ 4] = -w;    offset[ 5] = -w+1;
    offset[ 6] = w-1;   offset[ 7] = w;     offset[ 8] = w+1;
     
    for (int i = 0; i < 2; ++i) {
        ReadFibers(confPathname[i], vertPathname[i], &fibers, cw, ch);
        EuclideanParameterization(fibers, euclidean[i], cw*ch);
        delete[] fibers;
    }
 
    omp_set_num_threads(64);
 
    for (int z = 1; z < d - 1; ++z) {
        ReadFibers(confPathname[z+1], vertPathname[z+1], &fibers, cw, ch);
        EuclideanParameterization(fibers, euclidean[2], cw*ch);
        delete[] fibers;
 
        for (int y = 1; y < h - 1; ++y) {
#pragma omp parallel for            
            for (int x = 1; x < w - 1; ++x) {
                int index = y * w + x;
                Euclidean f[27];
 
                for (int i = 0; i < 9; ++i) {
                    memcpy(&(f[i]), &(euclidean[0][(index+offset[i])]), sizeof(Euclidean));
                    if (InnerProduct(f[i].seed.dir, f[0].seed.dir) < 0.0f)
                        ReverseFiber(&f[i]);
                    memcpy(&(f[i+9]), &(euclidean[1][(index+offset[i])]), sizeof(Euclidean));
                    if (InnerProduct(f[i+9].seed.dir, f[0].seed.dir) < 0.0f)
                        ReverseFiber(&f[i+9]);
                    memcpy(&(f[i+18]), &(euclidean[2][(index+offset[i])]), sizeof(Euclidean));
                    if (InnerProduct(f[i+18].seed.dir, f[0].seed.dir) < 0.0f)
                        ReverseFiber(&f[i+18]);
                }
 
                edgeStrength[z*w*h+index] = ComputeEdgeStrengthEuclidean(f, 27);
            }
        }
 
        memcpy(euclidean[0], euclidean[1], sizeof(Euclidean)*w*h);
        memcpy(euclidean[1], euclidean[2], sizeof(Euclidean)*w*h);
        memset(euclidean[2], 0, sizeof(Euclidean)*w*h);
 
        printf("z = %d\n", z);
    }

	/* save the result */
	CScalar *result = new CScalar();
	result->CreateScalar(w, h, d, edgeStrength);
	result->SaveScalarFile(dst_pathname);

	delete result;

	for (int z = 0; z < d; ++z) {
		SafeDeleteArray(confPathname[z]);
		SafeDeleteArray(vertPathname[z]);
	}
	SafeDeleteArray(confPathname);
	SafeDeleteArray(vertPathname);

	SafeDeleteArray(edgeStrength);

	SafeDeleteArray(euclidean[0]);
	SafeDeleteArray(euclidean[1]);
	SafeDeleteArray(euclidean[2]);
	SafeDeleteArray(euclidean);

	return 0;
}